# MNE-Python 和 YASA 结合的睡眠分析程序
# 功能：前5分钟提供实时指标，之后切换到YASA进行滑动窗口分析
# 优势：平衡了实时性、科学准确性与用户体验

import numpy as np
import mne
import yasa
import time
import csv
import pickle
import os
import math
from datetime import datetime
from NeuroPy3 import NeuroPy 

class SleepAnalyzer:
    def __init__(self, port="COM5", sfreq=512, log_step_sec=1, analysis_step_sec=3, window_sec=300):
        self.port = port
        self.sfreq = sfreq
        self.log_step_size = sfreq * log_step_sec  # 每秒记录日志
        self.analysis_step_size = sfreq * analysis_step_sec  # 每3秒分析一次
        self.window_size = sfreq * window_sec  # 5分钟窗口
        
        self.log_counter = 0  # 日志计数器
        self.analysis_counter = 0  # 分析计数器
        
        self.buffer = []
        self.poor_signal_buffer = []
        self.power_history = []  # 存储历史功率值用于平滑
        self.smooth_window = 5  # 5秒滑动窗口
        
        self.log_buffer = []  # 缓存日志数据
        self.log_buffer_size = 60  # 每60条记录写入一次
        
        self.neuropy = NeuroPy(port, 57600)
        self.log_file = "sleep_analysis_log.csv"
        
        self.previous_stage = "未知"
        
        self._init_log()
        
    def _init_log(self):
        """初始化CSV日志文件"""
        headers = ['时间戳', '数据长度(秒)', '睡眠阶段', 'Delta功率(dB)', 'Theta功率(dB)', 'Alpha功率(dB)', 'Beta功率(dB)', 'Gamma功率(dB)', '信号质量']
        with open(self.log_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(headers)
    
    def _data_callback(self, raw_value):
        """接收原始EEG数据的回调函数"""
        self.buffer.append(raw_value)
    
    def _poor_signal_callback(self, poor_signal_value):
        """信号质量回调函数"""
        self.poor_signal_buffer.append(poor_signal_value)

    def _flush_logs(self):
        """批量写入日志文件"""
        if self.log_buffer:
            with open(self.log_file, 'a', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                writer.writerows(self.log_buffer)
            self.log_buffer.clear()
            
    def _analyze_rule_based(self, data_row, previous_state):
        """
        基于规则的实时睡眠分析函数。
        """
        try:
            delta, theta, alpha, beta, gamma, quality = map(float, data_row)
        except (ValueError, IndexError):
            return "数据格式错误"

        if quality != 0.00:
            return "信号异常"

        power_values = [delta, theta, alpha, beta, gamma]
        power_array = np.array(power_values)
        highest_power = power_array.min()

        if delta == highest_power and delta > -30:
            return "深度睡眠"
        if (beta + gamma) > (delta + theta) and alpha > delta and previous_state != "深度睡眠":
            return "REM睡眠"
        if alpha == highest_power:
            return "放松"
        if previous_state == "放松" and delta < alpha:
            return "入睡"
        if theta < delta and theta < alpha:
            return "浅睡眠"
        if beta == highest_power or gamma == highest_power:
            return "清醒"

        return "状态未知"

    def _analyze_yasa_based(self):
        """
        基于YASA的滑动窗口睡眠分析函数。
        """
        if len(self.buffer) < self.window_size:
            return "数据不足"
        
        # 提取最新5分钟的数据作为分析窗口
        eeg_data = self.buffer[-self.window_size:]
        info = mne.create_info(['eeg'], sfreq=self.sfreq, ch_types='eeg')
        raw = mne.io.RawArray(np.array([eeg_data]), info)
        
        # YASA睡眠分期，使用0.05的Epoch长度以获得更细粒度的结果
        try:
            sls = yasa.sleep_staging(raw, eeg_name='eeg', sf_hypno=1/self.analysis_step_sec)
            # 返回最新一秒的分析结果
            return sls.iloc[-1]
        except Exception as e:
            return f"YASA分析出错: {e}"

    def _smooth_power_values(self, current_powers):
        """对功率值进行平滑处理"""
        self.power_history.append(current_powers)
        if len(self.power_history) > self.smooth_window:
            self.power_history.pop(0)
        
        return np.mean(self.power_history, axis=0)

    def _process_data(self):
        """处理缓冲区数据并进行分析"""
        try:
            buffer_length = len(self.buffer)
            
            # --- 每秒记录日志和进行分析 ---
            if buffer_length >= self.log_step_size * (self.log_counter + 1):
                self.log_counter += 1
                
                raw_delta = self.neuropy.delta
                raw_theta = self.neuropy.theta
                raw_alpha = (self.neuropy.lowAlpha + self.neuropy.highAlpha) / 2
                raw_beta = (self.neuropy.lowBeta + self.neuropy.highBeta) / 2
                raw_gamma = (self.neuropy.lowGamma + self.neuropy.midGamma) / 2
                signal_quality = self.neuropy.poorSignal
                
                if raw_delta <= 0 or raw_theta <= 0 or raw_alpha <= 0 or raw_beta <= 0 or raw_gamma <= 0:
                    print(f"第{self.log_counter}秒 - 数据异常，跳过分析")
                    return

                delta_power = 10 * math.log10(raw_delta)
                theta_power = 10 * math.log10(raw_theta)
                alpha_power = 10 * math.log10(raw_alpha)
                beta_power = 10 * math.log10(raw_beta)
                gamma_power = 10 * math.log10(raw_gamma)
                
                smoothed_powers = self._smooth_power_values([delta_power, theta_power, alpha_power, beta_power, gamma_power])
                data_row = [smoothed_powers[0], smoothed_powers[1], smoothed_powers[2], smoothed_powers[3], smoothed_powers[4], signal_quality]

                # --- 核心逻辑：动态切换分析引擎 ---
                if buffer_length < self.window_size:
                    # 5分钟前，使用规则引擎
                    current_stage = self._analyze_rule_based(data_row, self.previous_stage)
                    self.previous_stage = current_stage
                else:
                    # 5分钟后，使用YASA引擎
                    if buffer_length >= self.analysis_step_size * (self.analysis_counter + 1):
                        self.analysis_counter += 1
                        current_stage = self._analyze_yasa_based()
                        self.previous_stage = current_stage
                    else:
                        current_stage = self.previous_stage # 在两次YASA分析之间，保持状态不变

                print(f"第{self.log_counter}秒 - 状态: {current_stage}")
                self._log_result(1.0, current_stage, data_row[0], data_row[1], data_row[2], data_row[3], data_row[4], data_row[5])
        
        except Exception as e:
            print(f"数据处理出错: {e}")

    def _log_result(self, length, stage, delta, theta, alpha, beta, gamma, quality):
        """记录分析结果到日志"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        self.log_buffer.append([timestamp, length, stage, f"{delta:.6f}", f"{theta:.6f}", f"{alpha:.6f}", f"{beta:.6f}", f"{gamma:.6f}", f"{quality:.2f}"])
        
        if len(self.log_buffer) >= self.log_buffer_size:
            self._flush_logs()

    def start(self):
        """启动数据采集"""
        self.neuropy.setCallBack("rawValue", self._data_callback)
        self.neuropy.setCallBack("poorSignal", self._poor_signal_callback)
        self.neuropy.start()
        print(f"设备已连接到 {self.port}，开始采集数据...")
        
        while len(self.buffer) < self.log_step_size:
            time.sleep(0.1)
        
        try:
            while True:
                self._process_data()
                time.sleep(0.1)
        except KeyboardInterrupt:
            print("\n停止采集...")
        finally:
            self.neuropy.stop()
            self._flush_logs()
            print("设备已断开")

if __name__ == "__main__":
    analyzer = SleepAnalyzer(log_step_sec=1, analysis_step_sec=3, window_sec=300)
    analyzer.start()